# -*- coding: utf-8 -*-
# ----------------------------
# @Time    : 2024/1/4 11:17
# @Author  : changqingai
# @FileName: 分位数.py
# ----------------------------

from scipy.stats import norm
import torch
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']


def draw_single_plot(x, y, x_name='x', y_name='y', img_path=''):
    plt.figure(figsize=(5, 2.5))
    plt.plot(x, y)
    plt.xlabel(x_name)
    plt.ylabel(y_name)
    if img_path:
        plt.savefig(img_path)
    plt.grid()
    plt.show()


def draw_multi_plot(value_list, x_name, y_name, title, img_path):
    """
    :param value_list: [x, y, name]
    :return:
    """
    fig, ax = plt.subplots()  # 创建图实例
    for x, y, name in value_list:
        ax.plot(x, y, label=name)

    ax.set_xlabel(x_name)
    ax.set_ylabel(y_name)
    ax.set_title(title)
    ax.legend()
    plt.grid()

    # 是否保存图片
    if img_path:
        plt.savefig(img_path)
        print("成功保存图片")
    plt.show()
    print("success")


def get_quantile(offset=0.99, num_bins=16):
    quantile = norm.ppf(torch.linspace(1 - offset, offset, num_bins + 1)).tolist()  # 将[1-offset,offset]区间等分为16份
    tmp = [(quantile[1:][idx] + val) / 2 for idx, val in enumerate(quantile[:-1])]  # 计算分位数
    r_max, r_min = tmp[-1], tmp[0]
    S = (r_max - r_min)/(1 - (-1))
    Z = 1 - r_max / S
    Q = [x/S + Z for x in tmp]  # 分位数量化到[-1,1]
    print("quantile", quantile)
    print(Q)
    return quantile[: len(Q)], Q


if __name__ == "__main__":
    x, y = get_quantile()
    draw_single_plot(x, y, x_name="分位数", y_name="value", img_path="../../imgs/lora/quantitle.png")
